import pandas as pd
import numpy as np
import os

pd.options.display.max_columns = 200
pd.options.display.max_rows = 1000
pd.set_option('max_info_columns', 200)
pd.set_option('expand_frame_repr', False)
pd.set_option('expand_frame_repr', True)
pd.set_option('max_colwidth',1000)
pd.set_option('display.width',None)
pd.set_option('display.float_format', lambda x: '%.3f' % x)

#Set to directoy housing the excel files output/excel_beta_{}_in_32.xlsx ... 
os.chdir("")

### Read in the estimates

outcome = 'any_vadc'
ests = pd.read_excel("excel_beta_{}_in_32.xlsx".format(outcome), header=0).T.reset_index()
ests['samp'] = ests['index'].str.extract("samp([0-9])")
ests['bct_id'] = ests['index'].str.extract("bctXQFE([0-9]+)")
ests['est'] = ests[0]

# Reshape
ests = ests.loc[ests.bct_id.notnull(), ['samp','bct_id','est']]
ests = ests.pivot(index="bct_id", columns="samp", values="est").reset_index()
ests.columns = ['bct_id','months_dep',outcome]

# Add in the sampling errors
vcv = pd.read_excel("excel_varcov_{}_in_32.xlsx".format(outcome), header=0)
ests['var_months_dep'] = ests.apply(lambda x: vcv.loc[
                    vcv['Unnamed: 0'] == "samp1bctXQFE{}".format(x.bct_id),
                    "samp1bctXQFE{}".format(x.bct_id)].values[0], axis=1)
ests['var_{}'.format(outcome)] = ests.apply(lambda x: vcv.loc[
                    vcv['Unnamed: 0'] == "samp2bctXQFE{}".format(x.bct_id),
                    "samp2bctXQFE{}".format(x.bct_id)].values[0], axis=1)
ests['covvar_months_dep_{}'.format(outcome)] = ests.apply(lambda x: vcv.loc[
                    vcv['Unnamed: 0'] == "samp1bctXQFE{}".format(x.bct_id),
                    "samp2bctXQFE{}".format(x.bct_id)].values[0], axis=1)
allests = ests.copy()

for outcome in ['combat_death','noncom_death']:
    ests = pd.read_excel("excel_beta_{}_by_32.xlsx".format(outcome), header=0).T.reset_index()
    ests['samp'] = ests['index'].str.extract("samp([0-9])")
    ests['bct_id'] = ests['index'].str.extract("bctXQFE([0-9]+)")
    ests['est'] = ests[0]

    # Reshape
    ests = ests.loc[ests.bct_id.notnull(), ['samp','bct_id','est']].copy()
    ests = ests.pivot(index="bct_id", columns="samp", values="est").reset_index()
    ests.columns = ['bct_id','months_dep',outcome]
    
    allests = allests.merge(ests[['bct_id',outcome]], how='left', on='bct_id')

    # Add in the sampling errors
    vcv = pd.read_excel("excel_varcov_{}_by_32.xlsx".format(outcome), header=0)
    allests['var_{}'.format(outcome)] = ests.apply(lambda x: vcv.loc[
                        vcv['Unnamed: 0'] == "samp2bctXQFE{}".format(x.bct_id),
                        "samp2bctXQFE{}".format(x.bct_id)].values[0], axis=1)
    allests['covvar_months_dep_{}'.format(outcome)] = ests.apply(lambda x: vcv.loc[
                        vcv['Unnamed: 0'] == "samp1bctXQFE{}".format(x.bct_id),
                        "samp2bctXQFE{}".format(x.bct_id)].values[0], axis=1)

# Add in the N
samp = pd.read_stata('coefs_rf_qtr_new.dta')
allests['bct_id'] = pd.to_numeric(allests.bct_id)
allests = allests.merge(samp[['BCTxQTRid','count']], how='left', left_on='bct_id', right_on='BCTxQTRid')

# Save
allests.to_csv('ests_formatlab.csv'.format(outcome))

